#include <iostream>
#include <sys/select.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <cstring>

#define MAX_CLIENTS 5
#define BUFFER_SIZE 1024

int main() {
    int server_fd, new_socket;
    struct sockaddr_in address;
    int addrlen = sizeof(address);

    // Creating socket file descriptor
    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) == 0) {
        perror("socket failed");
        exit(EXIT_FAILURE);
    }

    // Setting up address structure
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(8080);

    // Binding the socket to the port
    if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("bind failed");
        close(server_fd);
        exit(EXIT_FAILURE);
    }

    // Listening for incoming connections
    if (listen(server_fd, MAX_CLIENTS) < 0) {
        perror("listen");
        close(server_fd);
        exit(EXIT_FAILURE);
    }

    fd_set readfds;
    int client_sockets[MAX_CLIENTS];
    for (int i = 0; i < MAX_CLIENTS; i++) {
        client_sockets[i] = -1;
    }

    while (true) {
        FD_ZERO(&readfds);
        FD_SET(server_fd, &readfds);
        int max_fd = server_fd;

        // Add connected clients to the set
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (client_sockets[i] != -1) {
                FD_SET(client_sockets[i], &readfds);
                if (client_sockets[i] > max_fd) {
                    max_fd = client_sockets[i];
                }
            }
        }

        // Wait for activity on any of the sockets
        if (select(max_fd + 1, &readfds, nullptr, nullptr, nullptr) > 0) {
            // Check for incoming connection
            if (FD_ISSET(server_fd, &readfds)) {
                if ((new_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) {
                    perror("accept");
                    continue;
                }

                // Add new socket to array
                for (int i = 0; i < MAX_CLIENTS; i++) {
                    if (client_sockets[i] == -1) {
                        client_sockets[i] = new_socket;
                        break;
                    }
                }
                std::cout << "New connection accepted" << std::endl;
            }

            // Check for data from clients
            for (int i = 0; i < MAX_CLIENTS; i++) {
                if (client_sockets[i] != -1 && FD_ISSET(client_sockets[i], &readfds)) {
                    char buffer[BUFFER_SIZE] = {0};
                    int valread = read(client_sockets[i], buffer, BUFFER_SIZE);
                    if (valread <= 0) {
                        // Client disconnected
                        std::cout << "Client disconnected" << std::endl;
                        close(client_sockets[i]);
                        client_sockets[i] = -1;
                    } else {
                        // Echo back the received message
                        std::cout << "Received: " << buffer << std::endl;
                        char send_buffer[] = "I am a server"; 
                        send(client_sockets[i], send_buffer, sizeof(send_buffer), 0);
                    }
                }
            }
        }
    }

    close(server_fd);
    return 0;
}